#!/usr/bin/env python3
"""
完整诊断 - 找出为什么查询返回结果太少
"""

import psycopg2
import requests
import os
import urllib3

urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)

PROXY_HOST = "172.26.134.84"
PROXY_PORT = 7890
PROXIES = {
    'http': f'http://{PROXY_HOST}:{PROXY_PORT}',
    'https': f'http://{PROXY_HOST}:{PROXY_PORT}'
}

def get_db_connection():
    try:
        return psycopg2.connect(dbname="vectortutorialdb")
    except:
        try:
            return psycopg2.connect(dbname="vectortutorialdb", host="/var/run/postgresql")
        except:
            return psycopg2.connect(
                dbname="vectortutorialdb",
                user="postgres",
                password="postgres",
                host="localhost"
            )

def generate_embedding(text):
    """生成嵌入向量"""
    try:
        with open("API_KEY.txt", "r") as f:
            api_key = f.read().strip()
    except:
        return None
    
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json"
    }
    
    data = {
        "input": text,
        "model": "text-embedding-3-small"
    }
    
    try:
        response = requests.post(
            "https://api.openai.com/v1/embeddings",
            headers=headers,
            json=data,
            proxies=PROXIES,
            verify=False,
            timeout=30
        )
        
        if response.status_code == 200:
            result = response.json()
            return result['data'][0]['embedding']
        else:
            return None
    except:
        return None

def full_diagnosis():
    print("=" * 80)
    print("🔬 完整诊断 - 查找问题根源")
    print("=" * 80)
    
    conn = get_db_connection()
    cur = conn.cursor()
    
    # 1. 查看所有数据
    print("\n1️⃣ 数据库中的所有记录:")
    cur.execute("SELECT id, content FROM embeddings ORDER BY id;")
    all_records = cur.fetchall()
    for r in all_records:
        print(f"   ID {r[0]}: {r[1]}")
    
    # 2. 检查 Apple 相关记录
    print("\n2️⃣ 包含 'Apple' 的记录:")
    apple_records = [r for r in all_records if 'Apple' in r[1]]
    if apple_records:
        for r in apple_records:
            print(f"   ✅ ID {r[0]}: {r[1]}")
    else:
        print("   ❌ 没有找到包含 'Apple' 的记录！")
    
    # 3. 检查 gaming laptop 相关记录
    print("\n3️⃣ 包含 'gaming' 或 'RGB' 的记录:")
    gaming_records = [r for r in all_records if 'gaming' in r[1].lower() or 'RGB' in r[1]]
    if gaming_records:
        for r in gaming_records:
            print(f"   ✅ ID {r[0]}: {r[1]}")
    else:
        print("   ❌ 没有找到游戏相关的记录！")
    
    # 4. 测试 Apple products 查询
    print("\n4️⃣ 测试 'Apple products' 查询:")
    query = "Apple products"
    print(f"   生成查询向量: {query}")
    query_vec = generate_embedding(query)
    
    if query_vec:
        q_str = '[' + ','.join(map(str, query_vec)) + ']'
        
        # 查询所有记录的相似度
        cur.execute("""
            SELECT 
                id,
                content,
                embedding <=> %s::vector AS distance,
                1 - (embedding <=> %s::vector) AS similarity
            FROM embeddings
            WHERE embedding IS NOT NULL
            ORDER BY distance
            LIMIT 12;
        """, (q_str, q_str))
        
        results = cur.fetchall()
        print(f"\n   所有记录按相似度排序:")
        for i, (id_val, content, dist, sim) in enumerate(results, 1):
            print(f"   {i}. ID{id_val} [距离:{dist:.4f}] [相似度:{sim:.4f}] {content[:60]}...")
    
    # 5. 测试 gaming laptop 查询
    print("\n5️⃣ 测试 'gaming laptop with RGB keyboard' 查询:")
    query2 = "gaming laptop with RGB keyboard"
    print(f"   生成查询向量: {query2}")
    query_vec2 = generate_embedding(query2)
    
    if query_vec2:
        q_str2 = '[' + ','.join(map(str, query_vec2)) + ']'
        
        cur.execute("""
            SELECT 
                id,
                content,
                embedding <=> %s::vector AS distance,
                1 - (embedding <=> %s::vector) AS similarity
            FROM embeddings
            WHERE embedding IS NOT NULL
            ORDER BY distance
            LIMIT 12;
        """, (q_str2, q_str2))
        
        results = cur.fetchall()
        print(f"\n   所有记录按相似度排序:")
        for i, (id_val, content, dist, sim) in enumerate(results, 1):
            print(f"   {i}. ID{id_val} [距离:{dist:.4f}] [相似度:{sim:.4f}] {content[:60]}...")
    
    # 6. 检查是否所有向量维度相同
    print("\n6️⃣ 检查向量一致性:")
    cur.execute("""
        SELECT id, length(embedding::text) 
        FROM embeddings 
        WHERE embedding IS NOT NULL
        ORDER BY id;
    """)
    lengths = cur.fetchall()
    
    length_set = set([l[1] for l in lengths])
    if len(length_set) == 1:
        print(f"   ✅ 所有向量长度一致")
    else:
        print(f"   ❌ 向量长度不一致！")
        for id_val, length in lengths:
            print(f"      ID {id_val}: {length} 字符")
    
    cur.close()
    conn.close()
    
    print("\n" + "=" * 80)
    print("诊断完成")
    print("=" * 80)

if __name__ == "__main__":
    full_diagnosis()